# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# --------------------------------------------------------
# Position embedding utils
# --------------------------------------------------------

import numpy as np

import torch

# --------------------------------------------------------
# 2D sine-cosine position embedding
# References:
# Transformer: https://github.com/tensorflow/models/blob/master/official/nlp/transformer/model_utils.py
# MoCo v3: https://github.com/facebookresearch/moco-v3
# --------------------------------------------------------
def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
    """
    grid_size: int of the grid height and width
    return:
    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
    """
    grid_h = np.arange(grid_size, dtype=np.float32)
    grid_w = np.arange(grid_size, dtype=np.float32)
    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
    grid = np.stack(grid, axis=0)

    grid = grid.reshape([2, 1, grid_size, grid_size])
    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
    if cls_token:
        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
    return pos_embed


def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
    """
    从网格生成二维的正余弦位置编码。

    参数：
        embed_dim (int): 位置编码的输出维度（必须是偶数）
        grid (tuple): 一个包含两个元素的元组，每个元素表示网格的大小（例如，高度 H 和宽度 W）

    返回：
        emb (ndarray): 位置编码，形状为 (H*W, D)，D为 embed_dim
    """
    # 确保 embed_dim 是偶数，以便将维度分为两部分分别处理
    assert embed_dim % 2 == 0

    # 使用 embed_dim 的一半来编码网格的高度 (H)
    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
    # 使用 embed_dim 的一半来编码网格的宽度 (W)
    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)

    # 将高度和宽度的编码在维度上拼接在一起，形成完整的二维位置编码
    emb = np.concatenate([emb_h, emb_w], axis=1)  # (H*W, D)
    return emb



def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
    """
    从一维位置生成正余弦位置编码。

    参数：
        embed_dim (int): 每个位置的输出维度
        pos (ndarray): 需要编码的位置列表，形状为 (M,)，M 为位置的数量

    返回：
        emb (ndarray): 对应位置的编码，形状为 (M, D)，D 为 embed_dim
    """
    # 确保 embed_dim 是偶数，以便将维度分为两部分分别处理
    assert embed_dim % 2 == 0

    # 创建 omega 数组，范围从 0 到 embed_dim / 2，用于正余弦函数的频率控制
    omega = np.arange(embed_dim // 2, dtype=float)
    # 将 omega 的值缩放到 [0, 1] 范围
    omega /= embed_dim / 2.
    # 设置频率的缩放因子为 1 / (10000^omega)
    omega = 1. / 10000**omega  # (D/2,)

    # 将位置数组展平，确保 pos 形状为 (M,)
    pos = pos.reshape(-1)  # (M,)
    # 使用外积计算位置编码的值，生成形状为 (M, D/2) 的矩阵
    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), 外积计算

    # 对编码值分别应用正弦和余弦函数
    emb_sin = np.sin(out)  # (M, D/2)
    emb_cos = np.cos(out)  # (M, D/2)

    # 将正弦和余弦部分拼接在一起，得到完整的正余弦位置编码
    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
    return emb




# --------------------------------------------------------
# Interpolate position embeddings for high-resolution
# References:
# DeiT: https://github.com/facebookresearch/deit
# --------------------------------------------------------
def interpolate_pos_embed(model, checkpoint_model):
    """
    对位置编码进行插值，使其适应不同模型的输入大小。

    参数：
        model (nn.Module): 当前模型，它需要从检查点加载位置编码。
        checkpoint_model (dict): 从检查点加载的模型参数字典，包含位置编码 (pos_embed)。

    过程：
        - 该函数将原始检查点中的位置编码插值到当前模型需要的大小。
        - 位置编码会通过双三次插值方法进行调整，使其与模型的 patch 数量相匹配。
        - 保留了额外的 token（例如：cls_token 和 dist_token）。
    """
    # 检查点模型中是否存在 'pos_embed'，即位置编码
    if 'pos_embed' in checkpoint_model:
        # 获取检查点中的位置编码
        pos_embed_checkpoint = checkpoint_model['pos_embed']
        # 获取原位置编码的维度（embedding size）
        embedding_size = pos_embed_checkpoint.shape[-1]
        # 当前模型的 patch 数量
        num_patches = model.patch_embed.num_patches
        # 计算额外的 token 数量（例如：cls_token 和 dist_token）
        num_extra_tokens = model.pos_embed.shape[-2] - num_patches
        # 获取检查点中的原始位置编码的高度（假设为正方形）
        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
        # 获取当前模型中位置编码的目标大小（patch 数量的平方根）
        new_size = int(num_patches ** 0.5)
        # 输出信息，显示位置编码的插值调整
        if orig_size != new_size:
            print("Position interpolate from %dx%d to %dx%d" % (orig_size, orig_size, new_size, new_size))
            # 提取额外的 token（如 cls_token 和 dist_token）
            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
            # 获取要进行插值调整的普通位置 token（排除额外的 token）
            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
            # 重塑位置 tokens 以适应插值操作，形状为 (batch_size, orig_size, orig_size, embedding_size)
            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
            # 使用双三次插值方法调整位置编码的大小
            pos_tokens = torch.nn.functional.interpolate(
                pos_tokens, size=(new_size, new_size), mode='bicubic', align_corners=False)
            # 恢复原来的维度顺序，并将结果展平
            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
            # 将额外的 token 和插值调整后的位置编码拼接起来
            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
            # 更新检查点中的位置编码
            checkpoint_model['pos_embed'] = new_pos_embed

